﻿

#include <string>

#include "mmdeploy/text_detector.hpp"
#include "mmdeploy/text_recognizer.hpp"
#include "utils/argparse.h"
#include "utils/mediaio.h"
#include "utils/visualize.h"


//DEFINE_string(device, "cpu", R"(Device name, e.g. "cpu", "cuda")");
//DEFINE_string(output, "text_ocr_output.jpg", "Output image path");

using mmdeploy::TextDetector;
using mmdeploy::TextRecognizer;

int main() {
    std::string ARGS_det_model = "./data/dbnet/";
    std::string ARGS_reg_model = "./data/satrn/";
    std::string ARGS_image = "C:/0N/junk/auto-recognition/0726-1/xj-auto-recognition/_backup/20240726145100/20240726145208836_0.png";
    std::string FLAGS_device = "cpu";
    std::string FLAGS_output = "text_ocr_output.jpg";

    cv::Mat img = cv::imread(ARGS_image);
    if (img.empty()) {
        fprintf(stderr, "failed to load image: %s\n", ARGS_image.c_str());
        return -1;
    }

    mmdeploy::Device device(FLAGS_device);
    TextDetector detector{ mmdeploy::Model(ARGS_det_model), device };
    TextRecognizer recognizer{ mmdeploy::Model(ARGS_reg_model), device };

    // apply the detector, the result is an array-like class holding references to
    // `mmdeploy_text_detection_t`, will be released automatically on destruction
    TextDetector::Result bboxes = detector.Apply(img);

    // apply recognizer, if no bboxes are provided, full image will be used; the result is an
    // array-like class holding references to `mmdeploy_text_recognition_t`, will be released
    // automatically on destruction
    TextRecognizer::Result texts = recognizer.Apply(img, { bboxes.begin(), bboxes.size() });

    // visualize results
    utils::Visualize v;
    auto sess = v.get_session(img);
    for (size_t i = 0; i < bboxes.size(); ++i) {
        mmdeploy_text_detection_t& bbox = bboxes[i];
        mmdeploy_text_recognition_t& text = texts[i];
        sess.add_text_det(bbox.bbox, bbox.score, text.text, text.length, i);
    }

    if (!FLAGS_output.empty()) {
        cv::imwrite(FLAGS_output, sess.get());
    }

    std::cout << "done" << std::endl;
    return 0;
}
